import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# AI模型配置
AI_CONFIG = {
    'model': os.getenv('MODEL_NAME', 'deepseek-chat'),
    'api_base': os.getenv('API_BASE_URL', 'https://api.deepseek.com/v1'),
    'api_key': os.getenv('DEEPSEEK_API_KEY'),
    'temperature': float(os.getenv('TEMPERATURE', '1')),
    'max_tokens': int(os.getenv('MAX_TOKENS', '4096')),
}

# OLLAMA模型配置
OLLAMA_CONFIG = {
    'model': os.getenv('OLLAMA_MODEL_NAME', 'deepseek-coder:latest'),
    'api_base': os.getenv('OLLAMA_API_BASE_URL', 'http://localhost:11434'),
    'temperature': float(os.getenv('OLLAMA_TEMPERATURE', '0.7')),
    'max_tokens': int(os.getenv('OLLAMA_MAX_TOKENS', '4096')),
}

# 提示模板
PROMPT_TEMPLATE = """基于以下数据库结构：

{table_description}

请将以下问题转换为SQL查询语句：

问题：{question}

请只返回SQL查询语句，不要包含任何其他解释。""" 